package pl.edu.agh.neuraleconomy.core.ta.indicator;

import java.util.ArrayList;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import pl.edu.agh.neuraleconomy.core.nn.CoreUtils;
import pl.edu.agh.neuraleconomy.model.exchange.Exchange;

// Relative Strength Index
public class RSICalculator implements ICalculator {
	private int period;

	public RSICalculator(int period) {
		this.period = period;
	}

	// Get Relative Strength Index
	public Map<Date, Double> calculate(List<Exchange> exchanges) {
		Map<Date, Double> rsiMap = new HashMap<Date, Double>();
		double [] rsi = calculate(CoreUtils.toDoubleArray(exchanges));
		for (int i = period; i < exchanges.size(); i++) {
			rsiMap.put(exchanges.get(i).getDate(), rsi[i-period]);
		}
		
		return rsiMap;
	}

	public double[] calculate(double[] data) {
		double [] res = new double[data.length - period];
		ArrayList<Double> change = new ArrayList<Double>();
		ArrayList<Double> gain = new ArrayList<Double>();
		ArrayList<Double> loss = new ArrayList<Double>();
		ArrayList<Double> averageGain = new ArrayList<Double>();
		ArrayList<Double> averageLoss = new ArrayList<Double>();
		double gainSum = 0;
		double lossSum = 0;

		// Initialize
		for (int i = 0; i < data.length; i++) {
			change.add(0.0);
			gain.add(0.0);
			loss.add(0.0);
			averageGain.add(0.0);
			averageLoss.add(0.0);
		}

		// Calculate "change", "gain" and "loss"
		for (int i = 1; i < data.length; i++) {
			change.set(i, data[i] - data[i - 1]);
			if (change.get(i) > 0) {
				gain.set(i, change.get(i));
				loss.set(i, 0.0);
			} else if (change.get(i) < 0) {
				loss.set(i, change.get(i) * (-1.0));
				gain.set(i, 0.0);
			}
		}

		// Calculate first "average gain" and "average loss"
		for (int i = 0; i < period; i++) {
			gainSum += gain.get(i);
		}
		for (int i = 0; i < loss.size(); i++) {
			lossSum += loss.get(i);
		}
		averageGain.set(period, gainSum / ((double) period));
		averageLoss.set(period, lossSum / ((double) period));

		// Calculate subsequent "average gain" and "average loss" values
		for (int i = period + 1; i < data.length; i++) {
			averageGain.set(i, (averageGain.get(i - 1) * ((double) (period - 1)) + gain.get(i)) / ((double) period));
			averageLoss.set(i, (averageLoss.get(i - 1) * ((double) (period - 1)) + loss.get(i)) / ((double) period));
		}

		// Calculate RSI
		double RS = 0; // Relative strength
		double RSIndex = 0; // Relative strength index
		for (int i = period; i < data.length; i++) {
			RS = averageGain.get(i) / averageLoss.get(i);
			RSIndex = 100 - 100 / (1 + RS);
			res[i-period] = RSIndex;
		}
		
		return res;
	}

	public int getInputLenForOutputLen(int outputLen) {
		return outputLen + period;
	}

}